#include "yolo5_pre.h"
#include <iostream>
#include <string>
#include <filesystem>

void printUsage(const char* programName) {
    std::cout << "数据集处理程序 - 支持NV12格式转换和图片填充 (C++版本)" << std::endl;
    std::cout << std::endl;
    std::cout << "用法: " << programName << " [选项]" << std::endl;
    std::cout << std::endl;
    std::cout << "必需参数:" << std::endl;
    std::cout << "  -s, --source <目录>    源目录路径（包含.raw和.png文件的目录）" << std::endl;
    std::cout << "  -d, --dest <目录>      目标目录路径（转换后的文件保存位置）" << std::endl;
    std::cout << std::endl;
    std::cout << "可选参数:" << std::endl;
    std::cout << "  -r, --rotation <角度>  顺时针旋转角度（默认0度，表示不旋转）" << std::endl;
    std::cout << "  -h, --help             显示此帮助信息" << std::endl;
    std::cout << std::endl;
    std::cout << "使用示例:" << std::endl;
    std::cout << "  " << programName << " -s /path/to/source -d /path/to/destination" << std::endl;
    std::cout << "  " << programName << " --source /path/to/source --dest /path/to/destination" << std::endl;
    std::cout << "  " << programName << " -s /path/to/source -d /path/to/destination -r 90" << std::endl;
    std::cout << std::endl;
    std::cout << "功能说明:" << std::endl;
    std::cout << "  - 自动检测并转换NV12格式的.raw文件直接为640x640的PNG格式" << std::endl;
    std::cout << "  - 将480x480的PNG图片填充到640x640（原图左上角，其余部分黑色）" << std::endl;
    std::cout << "  - 支持在填充前对图片进行顺时针旋转（-r参数）" << std::endl;
    std::cout << "  - 使用OpenCV进行图像转换，RGA2进行硬件加速处理" << std::endl;
    std::cout << "  - 零拷贝处理，直接在内存中完成转换和填充" << std::endl;
    std::cout << "  - 处理后的文件保存在指定的目标目录中" << std::endl;
}

int main(int argc, char* argv[]) {
    std::string sourceDir, destDir;
    double rotationAngle = 0.0;
    
    // 解析命令行参数
    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        
        if (arg == "-h" || arg == "--help") {
            printUsage(argv[0]);
            return 0;
        } else if (arg == "-s" || arg == "--source") {
            if (i + 1 < argc) {
                sourceDir = argv[++i];
            } else {
                std::cerr << "错误: -s/--source 参数需要指定目录路径" << std::endl;
                return 1;
            }
        } else if (arg == "-d" || arg == "--dest") {
            if (i + 1 < argc) {
                destDir = argv[++i];
            } else {
                std::cerr << "错误: -d/--dest 参数需要指定目录路径" << std::endl;
                return 1;
            }
        } else if (arg == "-r" || arg == "--rotation") {
            if (i + 1 < argc) {
                try {
                    rotationAngle = std::stod(argv[++i]);
                } catch (const std::exception& e) {
                    std::cerr << "错误: -r/--rotation 参数需要指定有效的数字" << std::endl;
                    return 1;
                }
            } else {
                std::cerr << "错误: -r/--rotation 参数需要指定旋转角度" << std::endl;
                return 1;
            }
        } else {
            std::cerr << "错误: 未知参数 '" << arg << "'" << std::endl;
            printUsage(argv[0]);
            return 1;
        }
    }
    
    // 检查必需参数
    if (sourceDir.empty()) {
        std::cerr << "错误: 必须指定源目录 (-s/--source)" << std::endl;
        printUsage(argv[0]);
        return 1;
    }
    
    if (destDir.empty()) {
        std::cerr << "错误: 必须指定目标目录 (-d/--dest)" << std::endl;
        printUsage(argv[0]);
        return 1;
    }
    
    // 验证源目录
    if (!Yolo5Pre::validateDirectory(sourceDir)) {
        std::cerr << "错误: 源目录 '" << sourceDir << "' 无效" << std::endl;
        return 1;
    }
    
    // 转换为绝对路径
    try {
        sourceDir = std::filesystem::absolute(sourceDir).string();
        destDir = std::filesystem::absolute(destDir).string();
    } catch (const std::exception& e) {
        std::cerr << "错误: 无法获取绝对路径: " << e.what() << std::endl;
        return 1;
    }
    
    std::cout << "=== YOLO5 数据集预处理程序 (C++版本) ===" << std::endl;
    std::cout << "使用OpenCV进行图像转换，RGA2进行硬件加速" << std::endl;
    std::cout << std::endl;
    
    try {
        // 创建Yolo5Pre实例并处理数据集
        Yolo5Pre processor;
        if (processor.processDataset(sourceDir, destDir, rotationAngle)) {
            std::cout << std::endl << "处理完成!" << std::endl;
            return 0;
        } else {
            std::cerr << std::endl << "处理失败!" << std::endl;
            return 1;
        }
    } catch (const std::exception& e) {
        std::cerr << "错误: 程序执行异常: " << e.what() << std::endl;
        return 1;
    }
}
